// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <cantor/parser/ast/base.h>
#include <cantor/parser/ast/func_expr.h>
#include <cantor/parser/ast/func_when_clause.h>
#include <cantor/parser/ast/func_ctor.h>

namespace cantor::parser {
    struct DqlNode : public DispatchExpr {
        DqlNode(pb::ast::ExprNodeType t) : DispatchExpr(t, pb::ast::ExprNodeClass::ENC_DQL) {
        }

        DqlNode(pb::ast::ExprNodeType t, pb::ast::ExprNodeClass c) : DispatchExpr(t, pb::ast::expr_class_merge(pb::ast::ExprNodeClass::ENC_DQL, c)) {
        }

        nlohmann::ordered_json to_json_impl() const override;
    };

    struct SubqueryExpr : public DqlNode {
        StmtNode *query_stmt = nullptr;

        SubqueryExpr() : DqlNode(pb::ast::ENT_SUB_QUERY_EXPR, pb::ast::ENC_SUBQUERY) {
            is_complex = true;
        }

        nlohmann::ordered_json to_json() const override;

        void print() const override {
            std::cout << this << std::endl;
        }

        void to_stream(std::ostream &os) const override;
    };

    struct ExistsSubqueryExpr : public DqlNode {
        SubqueryExpr *query_expr = nullptr;
        bool is_not = false;

        ExistsSubqueryExpr() : DqlNode(pb::ast::ENT_EXISTS_SUB_QUERY_EXPR, pb::ast::ENC_SUBQUERY) {
            is_complex = true;
        }

        nlohmann::ordered_json to_json() const override;

        void print() const override {
            std::cout << this << std::endl;
        }

        void to_stream(std::ostream &os) const override;
    };


    struct CompareSubqueryExpr : public DqlNode {
        DispatchExpr *left_expr = nullptr;
        SubqueryExpr *right_expr = nullptr;
        pb::ast::CompareType cmp_type = pb::ast::CMP_ANY;
        pb::ast::FuncType func_type = pb::ast::FT_COMMON;

        CompareSubqueryExpr() : DqlNode(pb::ast::ENT_COMPARE_SUBQUERY_EXPR, pb::ast::ENC_SUBQUERY) {
            is_complex = true;
        }

        nlohmann::ordered_json to_json() const override;

        void print() const override {
            std::cout << this << std::endl;
        }

        void to_stream(std::ostream &os) const override;

        const char *get_func_name() const;
    };

    struct CommonTableExpr : public DqlNode {
        SubqueryExpr *query_expr = nullptr;
        Vector<RangeVar *> column_names;
        String name;
        String option;

        CommonTableExpr() : DqlNode(pb::ast::ENT_COMMON_TABLE_EXPR) {
            is_complex = true;
            name = nullptr;
            option = nullptr;
        }

        nlohmann::ordered_json to_json() const override;


        void to_stream(std::ostream &os) const override;
    };


    struct Indirection : public DefiniteExpr {
        Indirection(pb::ast::IndirectionType t) : DefiniteExpr(pb::ast::ENT_INDIRECTION, pb::ast::ENC_DQL),
                                                indirection_type(t) {
        }

        nlohmann::ordered_json to_json_impl() const override {
            nlohmann::ordered_json j = DefiniteExpr::to_json_impl();
            j["indirection_type"] = indirection_type;
            return j;
        }

        pb::ast::IndirectionType indirection_type = pb::ast::IndirectionType::IND_NONE;
    };

    struct IndirectionIndex : public Indirection {
        IndirectionIndex() : Indirection(pb::ast::IndirectionType::IND_INDEX) {
        }

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override;

        DispatchExpr *index_expr = nullptr;
    };

    struct IndirectionSlice : public Indirection {
        IndirectionSlice() : Indirection(pb::ast::IndirectionType::IND_SLICE) {
        }

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override;

        DispatchExpr *slice_low{nullptr};
        DispatchExpr *slice_high{nullptr};
        DispatchExpr *slice_step{nullptr};
    };


    struct SelectField : public DefiniteExpr {
        DispatchExpr *expr = nullptr;
        String as_name;
        String org_name;
        bool wild_card{false};
        TableName wild_uri;

        SelectField() : DefiniteExpr(pb::ast::ENT_SELECT_FIELD, pb::ast::ENC_DQL) {
            as_name = nullptr;
            org_name = nullptr;
        }

        nlohmann::ordered_json to_json() const override;

        bool is_complex_node() override {
            if (expr != nullptr) {
                return expr->is_complex_node();
            } else {
                return false;
            }
        }


        void to_stream(std::ostream &os) const override;

        void find_placeholder(std::unordered_set<int> &placeholders) override {
            if (expr != nullptr) {
                expr->find_placeholder(placeholders);
            }
        }
    };

    struct ByItem : public DefiniteExpr {
        DispatchExpr *expr = nullptr;
        bool is_desc = false;

        ByItem() : DefiniteExpr(pb::ast::ENT_BY_ITEM, pb::ast::ENC_DQL) {
        }

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override {
            static const char *desc_str[] = {" ASC", " DESC"};
            os << expr << desc_str[is_desc];
        }
    };

    struct WhereClause : public DefiniteExpr {
        DispatchExpr *expr = nullptr;

        WhereClause() : DefiniteExpr(pb::ast::ENT_WHERE_CLAUSE, pb::ast::ENC_DQL) {
        }

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override {
            if (expr) {
                os << expr;
            }
        }
    };

    struct HavingClause : public DefiniteExpr {
        DispatchExpr *expr = nullptr;

        HavingClause() : DefiniteExpr(pb::ast::ENT_HAVING_CLAUSE, pb::ast::ENC_DQL) {
        }

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override {
            if (expr) {
                os << expr;
            }
        }
    };

    struct OrderByClause : public DefiniteExpr {
        Vector<ByItem *> items;

        OrderByClause() : DefiniteExpr(pb::ast::ENT_ORDER_BY, pb::ast::ENC_DQL) {
        }

        nlohmann::ordered_json to_json() const override;

        bool is_complex_node() override;

        void to_stream(std::ostream &os) const override {
            for (int i = 0; i < items.size(); i++) {
                os << " " << items[i];
                if (i != items.size() - 1) {
                    os << ",";
                }
            }
        }
    };

    struct GroupByClause : public DefiniteExpr {
        Vector<ByItem *> items;

        GroupByClause() : DefiniteExpr(pb::ast::ENT_GROUP_BY, pb::ast::ENC_DQL) {
        }

        nlohmann::ordered_json to_json() const override;

        bool is_complex_node() override {
            if (is_complex) {
                return true;
            }
            for (int i = 0; i < items.size(); i++) {
                if (items[i]->is_complex_node()) {
                    is_complex = true;
                    return true;
                }
            }
            return false;
        }

        void to_stream(std::ostream &os) const override {
            for (int i = 0; i < items.size(); i++) {
                os << " " << items[i];
                if (i != items.size() - 1) {
                    os << ",";
                }
            }
        }
    };

    struct LimitClause : public DefiniteExpr {
        DispatchExpr *offset;
        DispatchExpr *count;

        LimitClause() : DefiniteExpr(pb::ast::ENT_LIMIT_CLAUSE, pb::ast::ENC_DQL) {
            offset = nullptr;
            count = nullptr;
        }

        nlohmann::ordered_json to_json() const override;

        bool is_complex_node() override {
            if (is_complex) {
                return true;
            }
            if (offset->is_complex_node()) {
                is_complex = true;
                return true;
            }
            if (count->is_complex_node()) {
                is_complex = true;
                return true;
            }
            return false;
        }

        void to_stream(std::ostream &os) const override {
            os << " " << offset << ", " << count;
        }
    };


    struct WithClause : public DefiniteExpr {
        Vector<CommonTableExpr *> ctes;
        String                    option;

        WithClause() : DefiniteExpr(pb::ast::ENT_WITH_CLAUSE, pb::ast::ENC_DQL) {
            is_complex = true;
            option.set_null();
        }

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override {
            os << "WITH";
            for (int i = 0; i < ctes.size(); i++) {
                os << " ";
                ctes[i]->to_stream(os);
                if (i != ctes.size() - 1) {
                    os << ",";
                }
            }
        }
    };


    struct IndexHint : public DefiniteExpr {
        pb::ast::IndexHintType hint_type = pb::ast::IHT_HINT_NONE;
        pb::ast::IndexHintScope hint_scope = pb::ast::IHS_HINT_NONE;
        Vector<String> index_name_list;

        IndexHint() : DefiniteExpr(pb::ast::ENT_INDEX_HINT, pb::ast::ENC_DQL) {
        }

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override;
    };

    struct StmtNode;

    struct TableSource : public DefiniteExpr {
        TableName *table_name = nullptr;
        String as_name;
        StmtNode *derived_table = nullptr;
        Vector<IndexHint *> index_hints;
        Vector<String> partition_names;

        TableSource() : DefiniteExpr(pb::ast::ENT_TABLE_SOURCE, pb::ast::ENC_DQL) {
            as_name.set_null();
        }

        nlohmann::ordered_json to_json() const override;

        bool is_complex_node() override {
            if (derived_table != nullptr) {
                return true;
            }
            return false;
        }

        void to_stream(std::ostream &os) const override;
    };

    struct TableRef;

    struct JoinNode : public DefiniteExpr {
        TableRef *left = nullptr;
        TableRef *right = nullptr;
        pb::ast::JoinType join_type = pb::ast::JT_INNER_JOIN;
        DispatchExpr *expr = nullptr;
        Vector<RangeVar *> using_col;
        bool is_natural = false;
        bool is_straight = false;

        JoinNode() : DefiniteExpr(pb::ast::ENT_JOIN_NODE, pb::ast::ENC_DQL) {
            is_complex = true;
        }

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override;
    };


    struct TableRef {
        TableSource *source = nullptr;
        JoinNode *join_node = nullptr;

        nlohmann::ordered_json to_json() const {
            if (source != nullptr) {
                return source->to_json();
            }
            if (join_node != nullptr) {
                return join_node->to_json();
            }
            return nlohmann::ordered_json();
        }

        bool is_complex_node() const {
            if (source) {
                return source->is_complex_node();
            }
            if (join_node != nullptr) {
                return join_node->is_complex_node();
            }
            return false;
        }

        void to_stream(std::ostream &os) const {
            if (source) {
                source->to_stream(os);
            }
            if (join_node != nullptr) {
                join_node->to_stream(os);
            }
        }

        LocationInfo location;
    };

    static const char *priority_str[] = {"", " LOW_PRIORITY", " DELAYED", " HIGH_PRIORITY"};

    struct SelectStmtOpts : public DefiniteExpr {
        SelectStmtOpts() : DefiniteExpr(pb::ast::ENT_SELECT_STMT_OPT, pb::ast::ENC_DQL) {
        }

        bool distinct = false;
        bool sql_cache = false;
        bool calc_found_rows = false;
        bool straight_join = false;
        pb::ast::PriorityEnum priority{pb::ast::PE_NO_PRIORITY};

        nlohmann::ordered_json to_json() const override;

        void to_stream(std::ostream &os) const override;
    };
} // namespace cantor::parser
